#include <strmif.h>
#include "cstreamio.hpp"
#include "mkvstreamvideo.hpp"
#include "mkvstreamaudio.hpp"
#include "ebmlio.hpp"
#include <vfwmsgs.h>
#include <cassert>
#include <climits>
#include <ctime>
#include <limits>
#ifdef _DEBUG
#include "odbgstream.hpp"
using std::endl;
#endif

namespace MkvMux
{


CStreamIO::CStreamIO() :
    m_pVideo(0),
    m_pAudio(0)  //TODO: generalize to many audio streams
{
    //Seed the random number generator, which is needed
    //for creation of unique TrackUIDs.
    
    const time_t time_ = time(0);
    const unsigned seed = static_cast<unsigned>(time_);
    srand(seed);
}


CStreamIO::~CStreamIO()
{
    assert(!m_context.m_file.IsOpen());
    assert(m_pVideo == 0);
    assert(m_pAudio == 0);
}


void CStreamIO::SetVideo(MkvStreamVideo* p)
{
    assert(m_pVideo == 0);
    m_pVideo = p;
}
    

void CStreamIO::SetAudio(MkvStreamAudio* p)
{
    assert(m_pAudio == 0);
    m_pAudio = p;
}


void CStreamIO::Init(IStream* pStream)
{
    MkvStream::Context& ctx = m_context;
    EbmlIO::File& file = ctx.m_file;
    
    file.SetStream(pStream);
    
    ctx.m_max_timecode_ms = 0;  //keep track of duration

    if (file.IsOpen())   //pStream != NULL
    {        
        file.SetPosition(0);
        WriteEbmlHeader();
        InitSegment();
    }
}


void CStreamIO::Final()
{
    //TODO: we need to grant last wishes to the stream objects,
    //if we don't receive an EOS explicitly.
    
    MkvStream::Context& ctx = m_context;
    EbmlIO::File& file = ctx.m_file;
    
    if (file.IsOpen())
    {
        FinalSegment();
        file.SetStream(0);
    }

    delete m_pVideo;
    m_pVideo = 0;
    
    delete m_pAudio;
    m_pAudio = 0;
}


HRESULT CStreamIO::WriteVideo(IMediaSample* pSample)
{
    assert(pSample);
    assert(m_pVideo);
    
    return m_pVideo->Write(pSample);
}


HRESULT CStreamIO::WriteVideoEOS()
{
    assert(m_pVideo);    
    return m_pVideo->EndOfStream();
}



HRESULT CStreamIO::WriteAudio(IMediaSample* pSample)
{
    assert(pSample);
    assert(m_pAudio);
    return m_pAudio->Write(pSample);
}


HRESULT CStreamIO::WriteAudioEOS()
{
    assert(m_pAudio);    
    return m_pAudio->EndOfStream();
}


void CStreamIO::WriteEbmlHeader()
{
    EbmlIO::File& file = m_context.m_file;
    
    file.WriteID4(0x1A45DFA3);
    
    //Allocate 1 byte of storage for Ebml header size.
    const __int64 start_pos = file.SetPosition(1, STREAM_SEEK_CUR);

    //EBML Version
    
    file.WriteID2(0x4286);
    file.Write1UInt(1);
    file.Serialize1UInt(1);  //EBML Version = 1

    //EBML Read Version
    
    file.WriteID2(0x42F7);
    file.Write1UInt(1);
    file.Serialize1UInt(1);  //EBML Read Version = 1
    
    //EBML Max ID Length
    
    file.WriteID2(0x42F2);
    file.Write1UInt(1);
    file.Serialize1UInt(4);  //EBML Max ID Length = 4
    
    //EBML Max Size Length
    
    file.WriteID2(0x42F3);
    file.Write1UInt(1);
    file.Serialize1UInt(8);  //EBML Max Size Length = 8

    //Doc Type

    file.WriteID2(0x4282);    
    file.Write1String("matroska");
        
    //Doc Type Version
    
    file.WriteID2(0x4287);
    file.Write1UInt(1);
    file.Serialize1UInt(1);  //Doc Type Version = 1
    
    //Doc Type Read Version
    
    file.WriteID2(0x4285);
    file.Write1UInt(1);
    file.Serialize1UInt(1);  //Doc Type Read Version = 1
    
    const __int64 stop_pos = file.GetPosition();
    
    const __int64 size_ = stop_pos - start_pos;
    assert(size_ <= std::numeric_limits<BYTE>::max());
    
    const BYTE size = static_cast<BYTE>(size_);
    
    file.SetPosition(start_pos - 1);
    file.Write1UInt(size);
    
    file.SetPosition(stop_pos);
}


void CStreamIO::InitSegment()
{   
    MkvStream::Context& ctx = m_context;
    EbmlIO::File& file = ctx.m_file;
    
    ctx.m_segment_pos = file.GetPosition();
    
    file.WriteID4(0x18538067);  //SeekHead ID
    file.Write8UInt(0);         //will need to be filled in later
    
    InitSeekHead();  //Meta Seek
    InitInfo();      //Segment Info
    WriteTrack();
}


void CStreamIO::FinalSegment()
{
    MkvStream::Context& ctx = m_context;
    EbmlIO::File& f = ctx.m_file;
    
    m_cues_pos = f.GetPosition();  //end of clusters    
    ctx.WriteCues();
    
    const __int64 maxpos = f.GetPosition();    
    f.SetSize(maxpos);
    
    f.SetPosition(ctx.m_segment_pos);

    const ULONG id = f.ReadID4();
    assert(id == 0x18538067);
    id;
    
    const __int64 size = maxpos - ctx.m_segment_pos - 12;
    assert(size >= 0);
    
    f.Write8UInt(size);  //total size of the segment

    FinalSeekHead();    
    FinalInfo();    
    ctx.FinalClusters(m_cues_pos);
}


void CStreamIO::FinalInfo()
{
    MkvStream::Context& ctx = m_context;
    EbmlIO::File& f = ctx.m_file;

    f.SetPosition(m_duration_pos + 2 + 1);  //2 byte ID + 1 byte size
    
    //TODO: account for duration of last frame too    
    const float duration = static_cast<float>(ctx.m_max_timecode_ms);
    
    f.Serialize4Float(duration);
}


void CStreamIO::InitSeekHead()
{
    EbmlIO::File& file = m_context.m_file;
    
    m_seekhead_pos = file.GetPosition();
    
    //Let's try to write a cluster about every second.
    //We need to figure out about how may clusters we'll need.
    //The longest movie I think I've seen is Lawrence of Arabia,
    //which imdb.com is telling me is 227 min, or 13620 sec.
    //The director's cut of Cleopatra is a whopping 320 min, 
    //or 19200 sec.  So let's round up to 20000 as the max
    //number of seek entries.  (It's a pity we must do this,
    //but whatever.  We could make an optimization by using
    //IMediaSeeking to return the duration of the stream.)
    
    //The SeekID is 2 + 1 + 4 = 7 bytes.
    //The SeekPos is 2 + 1 + 8 = 11 bytes.
    //Total payload for a seek entry is 7 + 11 = 18 bytes.
    
    //The Seek entry is 2 + 1 + 18 = 21 bytes.
    //Total payload for SeekHead is 20000 * 21 = 420000 bytes.
    
    //A void element is 1 + 4 + len.
    //For now the Seek payload will be 1 + 4 + 420000 = 420005.
    
    const ULONG size = 420000;
    
    file.WriteID4(0x114D9B74);  //Seek Head
    file.Write4UInt(size + 5);  //5 = 1 byte ID + 4 byte size
    
    //The size of the SeekHead element is now fixed.  The
    //entries within payload part of the SeekHead will need 
    //adjusting later.  But the total size itself
    //is determined apriori (unlike the case for, say, the 
    //Segment element).
    
    //Note that we are giving outselves a budget of 20000
    //entries.  Once we spend our budget, we won't be able
    //be able to write any more level-1 elements to the file.
    //(Or we could write them, but not put mention them
    //in the SeekHead element.)
    
    //Note also that our SeekHead element will always contain
    //at least one Void element.  (In fact, it should contain
    //exactly one Void element.)  We have budgeted for this
    //when we specified the size of the SeekHead element.
    
    //TODO: add a CRC to this element.

    //We don't have the data yet, so we'll void it for now.

    file.WriteID1(0xEC);      //Void ID
    file.Write4UInt(size);    
    file.SetPosition(size, STREAM_SEEK_CUR);
}


void CStreamIO::FinalSeekHead()
{
    //TODO: it is permissible to write more than one seek head element.
    //That would mean we could have a small seek head up front,
    //comprising just 4 entries, for the info, track, cue, and other the seek head,
    //and we could have the other seek head have entries for just the clusters.
    //Alternatively, we could include the position of the first cluster
    //in the first seek head.
    
    //TODO: another possibility is that we've already budgeted a large
    //amount of data, we could see if we have enough space to put the 
    //cues up front.
    
    //TODO: splitter considerations: an async source can take a long
    //time to download, especially if this is a network stream.
    //We wouldn't have an index of keyframes until we've downloaded
    //the entire file (becuase the cues go at the end).  The benefit
    //of having cluster entries up front is that that would give us
    //a defacto index early.  Even without cues, you could get
    //rough indexing early.  You wouldn't get every keyframe (you
    //need the cues for that), but you can get something.  However,
    //if this is a network source, you can't really seek anyway,
    //since you don't yet have the frames!  However, we do wish
    //to minimize how much data needs to be downloaded before
    //playback can begin (even without the ability to seek),
    //so that would favor a smaller seekhead, and hence
    //fewer (but larger) clusters.  With 1ms resolution, a 
    //cluster can contain up to 32s of payload, so maybe a better
    //idea is to create clusters once every 10s or 20s or 30s,
    //instead of once every 1s as we do now.
    
    //Another option is to write the first cluster into the 
    //first seekhead (near the front of the file), and the remaining 
    //clusters into the second seekhead (near the back of the file).
    //That would minimize the amount of data required to begin
    //playback of the file.  (Really, who cares about the position
    //of clusters, except for the first one.  If you're really
    //interested in seeking, you're going to use the cues, not
    //the seekhead.)
    
    //Yet another option is to write two cues into the file.  The
    //first one containing a single entry, for the first keyframe
    //(at tc=0), and then have another cues at the end, with the
    //remaing keyframes.  Is this permitted?
    
    MkvStream::Context& ctx = m_context;
    EbmlIO::File& f = ctx.m_file;
    
    const __int64 start_pos = f.SetPosition(m_seekhead_pos + 8);
    
    WriteSeekEntry(0x1549A966, m_info_pos);   //Segment Info
    WriteSeekEntry(0x1654AE6B, m_track_pos);  //Track
 
    typedef MkvStream::clusters_t::const_iterator iter_t;
    
    MkvStream::clusters_t& cc = ctx.m_clusters;

    iter_t i = cc.begin();
    const iter_t j = cc.end();
    
    while (i != j)
    {
        const MkvStream::Cluster& c = *i++;        
        WriteSeekEntry(0x1F43B675, c.m_pos);
    }
    
    WriteSeekEntry(0x1C53BB6B, m_cues_pos);
    
    const __int64 stop_pos = f.GetPosition();
    
    const __int64 size_ = 420000 - (stop_pos - start_pos);
    assert(size_ >= 0);
    
    const ULONG size = static_cast<ULONG>(size_);
    
    f.WriteID1(0xEC);  //Void ID
    f.Write4UInt(size);
    //f.SetPosition(size, STREAM_SEEK_CUR);
}


void CStreamIO::WriteSeekEntry(ULONG id, __int64 pos_)
{
    //The SeekID is 2 + 1 + 4 = 7 bytes.
    //The SeekPos is 2 + 1 + 8 = 11 bytes.
    //Total payload for a seek entry is 7 + 11 = 18 bytes.
    
    //The Seek entry is 2 + 1 + 18 = 21 bytes.
    //Total payload for SeekHead is 20000 * 21 = 420000 bytes.
    
    MkvStream::Context& ctx = m_context;
    EbmlIO::File& f = ctx.m_file;

#ifdef _DEBUG    
    const __int64 start_pos = f.GetPosition();
#endif
    
    f.WriteID2(0x4DBB);  //Seek Entry ID
    f.Write1UInt(18);    //payload size of this Seek Entry
    
    f.WriteID2(0x53AB);  //SeekID ID
    f.Write1UInt(4);     //payload size is 4 bytes
    f.WriteID4(id);
    
    const __int64 pos = pos_ - ctx.m_segment_pos - 12;
    assert(pos >= 0);
    
    f.WriteID2(0x53AC);     //SeekPos ID
    f.Write1UInt(8);        //payload size is 8 bytes
    f.Serialize8UInt(pos);  //payload
    
#ifdef _DEBUG
    const __int64 stop_pos = f.GetPosition();
    assert((stop_pos - start_pos) == 21);
#endif
}


void CStreamIO::InitInfo()
{
    EbmlIO::File& file = m_context.m_file;
    
    m_info_pos = file.GetPosition();
    
    file.WriteID4(0x1549A966);  //Segment Info ID
    
    //allocate 1 byte of storage for size
    const __int64 pos = file.SetPosition(1, STREAM_SEEK_CUR);
    
    file.WriteID3(0x2AD7B1);       //TimeCodeScale ID
    file.Write1UInt(4);            //payload size
    file.Serialize4UInt(1000000);  //1ms resolution

    m_duration_pos = file.GetPosition();  //remember where duration is

    file.WriteID2(0x4489);         //Duration ID
    file.Write1UInt(4);            //payload size
    file.Serialize4Float(0.0);     //set value again during close
    
    //MuxingApp  //TODO
    //WritingApp  //TODO
    
    const __int64 newpos = file.GetPosition();

    const __int64 size_ = newpos - pos;
    assert(size_ <= std::numeric_limits<BYTE>::max());
    
    const BYTE size = static_cast<BYTE>(size_);
    
    file.SetPosition(pos - 1);
    file.Write1UInt(size);
    
    file.SetPosition(newpos);
}


void CStreamIO::WriteTrack()
{
    EbmlIO::File& file = m_context.m_file;

    m_track_pos = file.GetPosition();
    
    file.WriteID4(0x1654AE6B);  //Tracks element (level 1)
    
    //allocate 2 bytes of storage for size of Tracks element (level 1)
    const __int64 begin_pos = file.SetPosition(2, STREAM_SEEK_CUR);
    
    int tn = 0;
    
    if (m_pVideo)
        m_pVideo->WriteTrackEntry(++tn);
        
    if (m_pAudio)  //TODO: allow for multiple audio streams
        m_pAudio->WriteTrackEntry(++tn);
    
    const __int64 end_pos = file.GetPosition();
    
    const __int64 size_ = end_pos - begin_pos;
    assert(size_ <= USHRT_MAX);
    
    const USHORT size = static_cast<USHORT>(size_);
    
    file.SetPosition(begin_pos - 2);
    file.Write2UInt(size);
    
    file.SetPosition(end_pos);
}


} //end namespace MkvMux
